//
//  TriangleMesh.cpp
//  RayTracer
//
//  Created by 俞云康 on 3/25/16.
//  Copyright © 2016 yushroom. All rights reserved.
//

#include "TriangleMesh.hpp"
#include "Params.hpp"
#include "Timer.hpp"
#include "FileManager.hpp"
#include <fstream>

Triangle::Triangle(const std::shared_ptr<Transform>& transform)
    : Shape("", transform)
{
}

//Triangle::Triangle(uint32_t idx1, uint32_t idx2, uint32_t idx3)
//    : Shape("", mesh->transform)
//{
//    idx[0] = idx1;
//    idx[1] = idx2;
//    idx[2] = idx3;
//    //normal =
//}

bool Triangle::intersect(const Ray& r, HitInfo* hit) const
{
    auto& p1 = mesh->positions[idx[0]];
    auto& p2 = mesh->positions[idx[1]];
    auto& p3 = mesh->positions[idx[2]];
    //Ray r2 = mesh->transform.Ray2local(r);
	auto& r2 = r;
    
    // Transform triangle vertices to ray coordinate space
    
    float3 e1 = p2 - p1;
    float3 e2 = p3 - p1;
    float3 s1 = Cross(r2.dir, e2);
    float divisor = Dot(s1, e1);
    
    if (divisor == 0.)
        return false;
    float inv_divisor = 1.f / divisor;
    
    // Compute first barycentric coordinate
    float3 s = r2.o - p1;
    float b1 = Dot(s, s1) * inv_divisor;
    if (b1 < 0.f || b1 > 1.f)
        return false;
    
    // Compute second barycentric coordinate
    float3 s2 = Cross(s, e1);
    float b2 = Dot(r2.dir, s2) * inv_divisor;
    if (b2 < 0.f || b1 + b2 > 1.f)
        return false;
    
    float t = Dot(e2, s2) * inv_divisor;
    if (t < EPS || t > r2.maxt)
        return false;
    
    // Compute triangle partial derivatives
    float3 dpdu, dpdv;
    float2 uv[3];
    getUVs(uv);
    
    // Compute deltas for triangle partial derivatives
    float2 duv02 = uv[0] - uv[2];
    float2 duv12 = uv[1] - uv[2];
    Direction dp1 = p1 - p3, dp2 = p2 - p3;
    float det = duv02.x * duv12.y - duv02.y * duv12.x;
    if (det == 0) {
        CoordinateSystem(Normalize(Cross(e2, e1)), &dpdu, &dpdv);
    } else {
        float inv_det = 1.f / det;
        dpdu = ( duv12.y * dp1 - duv02.y * dp2) * inv_det;
        dpdv = (-duv12.x * dp1 + duv02.x * dp2) * inv_det;
    }
    
    // Interpolate $(u,v)$ triangle parametric coordinates
    float b0 = 1 - b1 - b2;
    float tu = b0*uv[0].x + b1*uv[1].x + b2*uv[2].x;
    float tv = b0*uv[0].y + b1*uv[1].y + b2*uv[2].y;
    
    r.maxt = t;
    hit->position = r.at(t);
    //hit->normal = mesh->transform.normal2world(normal);
    //hit->dpdu = mesh->transform.direction2world(dpdu);
    //hit->dpdv = mesh->transform.direction2world(dpdv);
	hit->normal = b0*mesh->normals[idx[0]]+b1*mesh->normals[idx[1]]+b2*mesh->normals[idx[2]];
    hit->normal.normalize_self();
	hit->dpdu = dpdu;
	hit->dpdv = dpdv;
    hit->texcoord.x = tu;
    hit->texcoord.y = tv;
    
    return true;
}

bool Triangle::intersect(const Ray& r) const
{
    auto& p1 = mesh->positions[idx[0]];
    auto& p2 = mesh->positions[idx[1]];
    auto& p3 = mesh->positions[idx[2]];
    //Ray r2 = mesh->transform.Ray2local(r);
	auto& r2 = r;
    
    // Transform triangle vertices to ray coordinate space
    
    float3 e1 = p2 - p1;
    float3 e2 = p3 - p1;
    float3 s1 = Cross(r2.dir, e2);
    float divisor = Dot(s1, e1);
    
    if (divisor == 0.)
        return false;
    float inv_divisor = 1.f / divisor;
    
    // Compute first barycentric coordinate
    float3 s = r2.o - p1;
    float b1 = Dot(s, s1) * inv_divisor;
    if (b1 < 0.f || b1 > 1.f)
        return false;
    
    // Compute second barycentric coordinate
    float3 s2 = Cross(s, e1);
    float b2 = Dot(r2.dir, s2) * inv_divisor;
    if (b2 < 0.f || b1 + b2 > 1.f)
        return false;
    
    float t = Dot(e2, s2) * inv_divisor;
    if (t < EPS || t > r2.maxt)
        return false;
    
    return true;
}


void Triangle::getUVs(float2 uv[3]) const
{
    if (mesh->uvs.empty()) {
        uv[0] = float2(0, 0);
        uv[1] = float2(1, 0);
        uv[2] = float2(1, 1);
    } else {
        uv[0] = mesh->uvs[idx[0]];
        uv[1] = mesh->uvs[idx[1]];
        uv[2] = mesh->uvs[idx[2]];
    }
}

float Triangle::area() const {
	const auto &p1 = mesh->positions[idx[0]];
	const auto &p2 = mesh->positions[idx[1]];
	const auto &p3 = mesh->positions[idx[2]];
	return 0.5f * Cross(p2 - p1, p3 - p1).length();
}

Position Triangle::sample(const Sample& sample, Normal* shape_normal) const
{
	float b1, b2;
	UniformSampleTriangle(sample.x, sample.y, &b1, &b2);
	const auto &p1 = mesh->positions[idx[0]];
	const auto &p2 = mesh->positions[idx[1]];
	const auto &p3 = mesh->positions[idx[2]];
	auto p = b1*p1 + b2*p2 + (1.f - b1 - b2)*p3;
	*shape_normal = Normal(Cross(p2 - p1, p3 - p1));
	shape_normal->normalize_self();
	return p;
}

TriangleMesh::TriangleMesh(const string& name,
             const std::shared_ptr<Transform>& transform,
             const string& path)
: name(name), transform(transform)
{
    loadFromObj(path);

	if (!transform->isIdentity()) {
		for (int i = 0; i < n_vertices; ++i) {
			positions[i] = transform->point2world(positions[i]);
			normals[i] = transform->normal2world(normals[i]);
		}
	}

    triangles.resize(n_triangles);
    for (int i = 0 ; i < n_triangles; ++i) {
        //triangles[i] = Triangle(indices[i*3], indices[i*3+1], indices[i*3+2]);
        triangles[i] = make_shared<Triangle>(transform);
        triangles[i]->idx[0] = indices[i*3];
        triangles[i]->idx[1] = indices[i*3+1];
        triangles[i]->idx[2] = indices[i*3+2];
        triangles[i]->mesh = this;
        auto& p1 = positions[indices[i*3]];
        auto& p2 = positions[indices[i*3+1]];
        auto& p3 = positions[indices[i*3+2]];
        //triangles[i]->normal = Normalize(Cross(p2-p1, p3-p1));
		triangles[i]->boundingBox.expandBy(p1);
		triangles[i]->boundingBox.expandBy(p2);
		triangles[i]->boundingBox.expandBy(p3);
//        boundingBox.expandBy(transform.point2world(p1));
//        boundingBox.expandBy(transform.point2world(p2));
//        boundingBox.expandBy(transform.point2world(p3));
    }
}

shared_ptr<TriangleMesh> TriangleMesh::create(const Params& params) {
    string name = params.getID();
    auto transform = Transform::create(params.child("transform"));
    string path = FileManager::getWorkingDir() + "/" + params.getString("filename");
    return make_shared<TriangleMesh>(name, transform, path);
}

//bool TriangleMesh::intersect(const Ray& r, HitInfo* hit) const {
//    bool ist = false;
//    Ray r2 = transform.Ray2local(r);
//    for (auto& t : triangles) {
//        if (t.intersect(r2, hit)) {
//            ist = true;
//            r.maxt = r2.maxt;
//        }
//    }
//    if (ist) {
//        r.maxt = r2.maxt;
//        hit->normal = transform.normal2world(hit->normal);
//        hit->position = transform.point2world(hit->position);
//        hit->dpdu = transform.direction2world(hit->dpdu);
//        hit->dpdv = transform.direction2world(hit->dpdv);
//    }
//    return ist;
//}
//
//bool TriangleMesh::intersect(const Ray& r) const {
//    bool ist = false;
//    for (auto& t : triangles) {
//        if (t.intersect(r)) {
//            ist = true;
//        }
//    }
//    return ist;
//}

namespace ModelProcessing {
    struct HalfEdge;
    
    struct Face {
        HalfEdge*	half_edge;
        int id;
    };
    
    struct Vertex {
        HalfEdge*	outgoing_halfedge = nullptr;
    };
    
    struct HalfEdge {
        HalfEdge*	next_half_edge;
        HalfEdge*	opposite_half_edge	= nullptr;
        Face*		face;
        Vertex*		to_vertex;	// the vertex it points to
    };
    
    class Mesh {
    private:
        vector<HalfEdge>        _half_edges;
        vector<Vertex>          _vertices;
        vector<Face>            _faces;
        
        vector<float3>          _positions;
        vector<unsigned int>    _indices;
        vector<float2>          _uvs;
        vector<unsigned int>    _uv_indices;
        
    public:
        // triangles: 3 vertex ids for each face
        Mesh(const vector<float3>&          points,
             const vector<unsigned int>&    triangles,
             const vector<float2>&          uvs,
             const vector<unsigned int>&    uv_indices);
        
        const vector<float3>& optimized_positions() const {
            return _positions;
        }
        
        const vector<unsigned int>& optimized_indices() const {
            return _indices;
        }
        
        const vector<float2>& optimized_uvs() const {
            return _uvs;
        }
        
        const vector<unsigned int>& optimized_uv_indices() const {
            return _uv_indices;
        }
        
        void generate_normals_and_tangents(vector<float3>* normals, vector<float3>* tangents) const;
    };
    
    template<typename T>
    int binary_search_pos(const vector<T>& vec, const T& v) {
        int left = 0, right = (int)vec.size(), mid;
        while (left <= right) {
            mid = (left+right)/2;
            if (v < vec[mid]) {
                right = mid - 1;
            } else if (v == vec[mid]) {
                return mid;
            } else {
                left = mid + 1;
            }
        }
        error("not found!");
        return -1;
    };
    
    Mesh::Mesh(const vector<float3>&        points,
               const vector<unsigned int>&  triangles,
               const vector<float2>&        uvs,
               const vector<unsigned int>&  uv_indices)
    : _positions(points), _indices(triangles), _uvs(uvs), _uv_indices(uv_indices)
    {
        auto face_count = triangles.size() / 3;
        _faces.resize(face_count);
        _half_edges.resize(3 * face_count);
        
        //cout << "old size = " << _positions.size() << endl;
        sort(_positions.begin(), _positions.end());
        auto last = std::unique(_positions.begin(), _positions.end());
        _positions.erase(last, _positions.end());
        //cout << "new size = " << _positions.size() << endl;
        if (points.size() != _positions.size()) {
            info("optimize position: remove %d positions.\n", points.size() - _positions.size());
        }
        
        for (auto& i : _indices) {
            i = binary_search_pos(_positions, points[i]);
        }
        
        _vertices.resize(_positions.size());
        
        {
            //cout << endl << uvs.size() << endl;
            std::sort(_uvs.begin(), _uvs.end());
            auto last = std::unique(_uvs.begin(), _uvs.end());
            _uvs.erase(last, _uvs.end());
            if (uvs.size() != _uvs.size()) {
                info("optimize uv: remove %d uvs.\n", uvs.size() - _uvs.size());
            }
            
            for (auto& i : _uv_indices) {
                i = binary_search_pos(_uvs, uvs[i]);
            }
        }
        
        map<pair<int, int>, HalfEdge*> edges;
        
        for (int i = 0; i < face_count; i++)
        {
            Face* f = &_faces[i];
            f->id = i;
            f->half_edge = &_half_edges[i * 3];
            for (int j = 0; j < 3; j++) {
                int id1 = _indices[i*3+j];
                int id2 = _indices[i*3+(j + 1) % 3];
                HalfEdge*   he      = &_half_edges[i * 3 + j];
                Vertex*     v       = &_vertices[id1];
                Vertex*     next_v  = &_vertices[id2];
                
                auto search = edges.find(make_pair(id2, id1));
                if (search != edges.end()) {
                    auto o_he = search->second;
                    he->opposite_half_edge = o_he;
                    o_he->opposite_half_edge = he;
                    edges.erase(search);
                }
                else {
                    edges[make_pair(id1, id2)] = he;
                }
                
                if (v->outgoing_halfedge == nullptr)
                    v->outgoing_halfedge = he;
                he->next_half_edge = &_half_edges[i * 3 + (j + 1) % 3];
                he->face = f;
                he->to_vertex = next_v;
            }
        }
        
        info("number of bound edges: %lu\n", edges.size());
        //Assert(edges.size() == 0);
    }
    
    void Mesh::generate_normals_and_tangents(vector<float3>* normals, vector<float3>* tangents) const
    {
        bool gen_tangents = (tangents != nullptr);
        bool gen_normals = gen_tangents || (normals != nullptr);
        assert(gen_normals || gen_tangents);
        if (normals != nullptr) {
            //normals->clear();
            //normals->reserve(_positions.size());
            normals->resize(_positions.size());
        }
        if (gen_tangents) {
            //tangents->clear();
            //tangents->reserve(_positions.size());
            tangents->resize(_positions.size());
        }
        
        // Assimp
        for (int i = 0; i < _faces.size(); ++i) {
            auto& p0 = _positions[_indices[i * 3]];
            auto& p1 = _positions[_indices[i * 3 + 1]];
            auto& p2 = _positions[_indices[i * 3 + 2]];
            float3 v = p1 - p0;
            float3 w = p2 - p0;
            float3 normal = Cross(v, w);
            float3 tangent, bitangent;
            if (gen_tangents) {
                auto& uv0 = _uvs[_indices[i * 3]];
                auto& uv1 = _uvs[_indices[i * 3 + 1]];
                auto& uv2 = _uvs[_indices[i * 3 + 2]];
                float sx = uv1.x - uv0.x;
                float sy = uv1.y - uv0.y;
                float tx = uv2.x - uv0.x;
                float ty = uv2.y - uv0.y;
                float dirCorrection = (tx * sy - ty * sx) < 0.0f ? -1.0f : 1.0f;
                // when t1, t2, t3 in same position in UV space, just use default UV direction.
                if (0 == sx && 0 == sy && 0 == tx && 0 == ty) {
                    sx = 0.0f; sy = 1.0f;
                    tx = 1.0f; ty = 0.0f;
                }
                
                tangent.x = (w.x * sy - v.x * ty) * dirCorrection;
                tangent.y = (w.y * sy - v.y * ty) * dirCorrection;
                tangent.z = (w.z * sy - v.z * ty) * dirCorrection;
                //                bitangent.x = (w.x * sx - v.x * tx) * dirCorrection;
                //                bitangent.y = (w.y * sx - v.y * tx) * dirCorrection;
                //                bitangent.z = (w.z * sx - v.z * tx) * dirCorrection;
            }
            for (int j = 0; j < 3; ++j) {
                int idx = _indices[i*3 + j];
                if (normals != nullptr)
                    (*normals)[idx] += normal;
                if (tangents != nullptr) {
                    //auto& p = _positions[idx];
                    float3 localTangent = tangent - normal * (Dot(tangent, normal));
                    //float3 localBitangent = bitangent - normal * (Dot(bitangent, normal));
                    //localTangent.NormalizeSelf();
                    //localBitangent.NormalizeSelf();
                    (*tangents)[idx] += localTangent;
                }
            }
        }
        
        for (int i = 0; i < _positions.size(); ++i) {
            //int idx = _indices[i];
            if (normals != nullptr) {
                (*normals)[i].normalize_self();
            }
            if (tangents != nullptr) {
                (*tangents)[i].normalize_self();
            }
        }
    }
}

void split_string(const string & str, int *vertex_id, int *uv_id, int *normal_id)
{
    int count = 0;
    for (unsigned int i = 0; i < str.length(); i++)
    {
        if (str[i] == '/') count++;
    }
    *vertex_id = *uv_id = *normal_id = 0;
    if (count == 0)
        sscanf(str.c_str(), "%d", vertex_id);
    else if (count == 1)
        sscanf(str.c_str(), "%d/%d", vertex_id, uv_id);
    else if (count == 2)
        sscanf(str.c_str(), "%d/%d/%d", vertex_id, uv_id, normal_id);
    else
        error("vertex not in wanted format in load_obj: not triangle\n");
    
    ////http://en.cppreference.com/w/cpp/language/operator_precedence
    //// the precedence of -- is prior than *
    //--(*vertex_id);
    //--(*uv_id);
    //--(*normal_id);
}

void TriangleMesh::loadFromObj(const string &path)
{
    Timer timer2("Load Model");
    timer2.begin();
    
    //Model m;
    cout << "Loading model " << path << endl;
    
    ifstream fin(path);
    if (!fin.good()) {
        //cout << "file not exists: " << path << endl;
        error("file not exists: %s\n", path.c_str());
        return;
    }
    
    int n_triangles = 0;
    int n_vertices = 0;
    vector<float3> positions;
    vector<float3> normals;
    vector<float3> tangents;
    vector<float2> uvs;
    vector<unsigned int> p_index;
    vector<unsigned int> uv_index;
    vector<unsigned int> n_index;
    
    char buffer[256];
    //char str[256];
    float f1, f2, f3;
    char face[3][32];
    int line_number = 0;
    
    while (!fin.eof())
    {
        fin.getline(buffer, 255);
        buffer[255] = '\0';
        //sscanf(buffer, "%s", str);
        //info(buffer);
        line_number++;
        
        // vertex
        if ((buffer[0] == 'v') && (buffer[1] == ' ' || buffer[1] == 32)) {
            if (sscanf(buffer, "v %f %f %f", &f1, &f2, &f3) == 3) {
                positions.push_back(float3(f1, f2, f3));
                //info("%f %f %f\n", f1, f2, f3);
                n_vertices++;
            }
            else {
                error("vertex not in wanted format in load_obj at line: %d\n", line_number);
                //return m;
            }
        }
        else if(buffer[0] == 'v' && (buffer[1] == 't')) {
            if (sscanf(buffer, "vt %f %f", &f1, &f2) == 2) {
                uvs.push_back(float2(f1, f2));
            }
            else {
                error("vertex not in wanted format in load_obj at line: %d\n", line_number);
                //return m;
            }
        }
        else if (buffer[0] == 'v' && (buffer[1] == 'n')) {
            if (sscanf(buffer, "vn %f %f %f", &f1, &f2, &f3) == 3) {
                normals.push_back(float3(f1, f2, f3));
            }
            else {
                error("vertex not in wanted format in load_obj at line: %d\n", line_number);
                //return m;
            }
        }
        else if(buffer[0] == 'f' && buffer[1] == ' ') {
            if (sscanf(buffer, "f %s %s %s", (char*)face, (char*)(face+1), (char*)(face+2)) == 3) {
                int v, u, n;
                //info("%s %s %s\n", face, face + 1, face + 2);
                for (int i = 0; i < 3; i++)
                {
                    split_string(face[i], &v, &u, &n);
                    //--v, --u, --n;
                    if (v < 0) v = (int)positions.size() + v + 1;
                    if (u < 0) u = (int)uvs.size() + u + 1;
                    if (n < 0) n = (int)normals.size() + n + 1;
                    if (v < 0 || u < 0 || n < 0) {
                        error("    face index error [load_obj] at line: %d\n", line_number);
                        //return m;
                    }
                    p_index.push_back(v-1);
                    if (u != 0)
                        uv_index.push_back(u-1);
                    if (n != 0)
                        n_index.push_back(n-1);
                }
                n_triangles++;
            }
            else {
                error("vertex not in wanted format in load_obj at line: %d\n", line_number);
                //return m;
            }
        }
    }
    
    Assert(positions.size() == n_vertices);
    Assert(p_index.size() == n_triangles * 3);
    
    info("Model loading finished. n_triangles = %d, n_vertices = %d.\n", n_triangles, n_vertices);
    
    if (n_index.size() != n_triangles * 3) {
        //error("some vertices do not have normal!\n");
        info("some vertices do not have normal. Generate normals...\n");
        ModelProcessing::Mesh mesh(positions, p_index, uvs, uv_index);
        positions = mesh.optimized_positions();
        p_index = mesh.optimized_indices();
        uvs = mesh.optimized_uvs();
        uv_index = mesh.optimized_uv_indices();
        mesh.generate_normals_and_tangents(&normals, nullptr);
        // reassign normal index
        n_index.clear();
        n_index.reserve(n_triangles * 3);
        for (int i = 0; i < n_triangles * 3; ++i) {
            n_index.push_back(p_index[i]);
        }
        //return nullptr;
    }
    
    bool use_uv = true;
    
    //if (!normals.empty()) {
    //Assert(n_index.size() == n_triangles * 3);
    if (n_index.size() != n_triangles * 3) {
        error("some vertices do not have normal!\n");
        //return m;
    }
    //}
    //if (!uvs.empty()) {
    //Assert(uv_index.size() == n_triangles * 3);
    if (uv_index.size() != n_triangles * 3) {
        warning("some vertices do not have uv!\n");
        use_uv = false;
        //return m;
    }
    //}
    
    n_vertices = (int)positions.size();
    n_triangles = (int)p_index.size() / 3;
    
    info("ModelProcessing\n");
    
    
    vector<float3> new_points;
    vector<float2> new_uvs;
    vector<float3> new_normals;
    vector<int> uv_id_at_p(n_vertices, -1);
    
    //int same_uv_count = 0;
    
    new_uvs.resize(n_vertices);
    new_normals.resize(n_vertices);
    uv_id_at_p.resize(n_vertices);
    if (use_uv) {
        for (int i = 0; i < n_triangles; i++) {
            int v = i * 3;
            for (int j = 0; j < 3; ++j) {
                int id_at_p_array = p_index[v];
                if (uv_id_at_p[id_at_p_array] == -1) {
                    uv_id_at_p[id_at_p_array] = uv_index[v];
                    new_uvs[id_at_p_array] = uvs[uv_index[v]];
                    new_normals[id_at_p_array] = normals[n_index[v]];
                }
                else if (uv_id_at_p[id_at_p_array] != uv_index[v]) {	// this vertex has other uv already
                    //                auto uv1 = uvs[uv_index[uv_id_at_p[id_at_p_array]]];
                    //                auto uv2 = uvs[uv_index[v]];
                    //                if (fabsf(uv1.x - uv2.x) < 1e-2f && fabsf(uv1.y - uv2.y) < 1e-2f) {
                    //                    same_uv_count ++;
                    //                }
                    int new_id = (int)positions.size();
                    positions.push_back(positions[p_index[v]]);				// add a new vertex
                    new_uvs.push_back(uvs[uv_index[v]]);
                    new_normals.push_back(normals[n_index[v]]);
                    p_index[v] = new_id;
                }
                v++;
            }
        }
    } else {
        new_uvs = uvs;
        new_normals = normals;
    }
    
    //info("same uv count = %d", same_uv_count);
    
    info("Model processing finished. Add %d new vertices. %lu vertices in all\n", int(positions.size()) - n_vertices, positions.size());
    n_vertices = (int)positions.size();
    
    this->n_triangles = n_triangles;
    this->n_vertices = n_vertices;
    this->indices = std::move(p_index);
    this->positions = std::move(positions);
    this->normals = std::move(new_normals);
    this->tangents = std::move(tangents);
    this->uvs = std::move(new_uvs);
    
    timer2.end();
    timer2.print();
}